#include <sys/time.h>
#include "maxpool_2x2.h"
#include "utils.h"

int main()
{
    // Test parameters for maxpool 2x2
    constexpr uint32_t N_CH = 8;     // Number of channels
    constexpr uint32_t N_IH = 8;     // Input height
    constexpr uint32_t N_IW = 8;     // Input width
    constexpr uint32_t N_BATCH = 1;  // Batch size
    constexpr uint32_t P_CH = 4;     // Parallel channels
    constexpr uint32_t FOLD = N_CH / P_CH;
    static_assert(N_CH % P_CH == 0, "N_CH must be divisible by P_CH");
    
    constexpr uint32_t N_OH = N_IH / 2;  // Output height after 2x2 maxpool
    constexpr uint32_t N_OW = N_IW / 2;  // Output width after 2x2 maxpool
    constexpr uint32_t N_IM = N_IH * N_IW * N_CH;  // Input size
    constexpr uint32_t N_OM = N_OH * N_OW * N_CH;  // Output size

    std::array<uint32_t, 4> ifm_shape = {N_BATCH, N_IH, N_IW, N_CH};
    std::array<uint32_t, 4> ofm_shape = {N_BATCH, N_OH, N_OW, N_CH};

    // Create input streams
    data_stream<BIT_ACTV> s_input("s_input");
    data_stream<BIT_ACTV> s_input_golden("s_input_golden");
    
    // Generate test input data
    gen_ifm(s_input, ifm_shape);
    gen_ifm(s_input_golden, ifm_shape);

    // Expand input width for parallel processing
    data_stream<BIT_ACTV * P_CH> s_input_e("s_input_e");
    expandWidth<BIT_ACTV, P_CH * BIT_ACTV, N_IM>(s_input, s_input_e);

    // Create output stream for the optimized version
    data_stream<P_CH * BIT_ACTV> s_output_e("s_output_e");
    
    // Run the optimized maxpool2x2 function
    maxpool2x2<P_CH, N_CH, N_IH, N_IW, N_BATCH, BIT_ACTV>(
        s_input_e, s_output_e
    );

    // Reduce output width back to single channel
    data_stream<BIT_ACTV> s_output("s_output");
    reduceWidth<BIT_ACTV * P_CH, BIT_ACTV, N_OM / P_CH>(s_output_e, s_output);

    // Run the golden reference function
    data_stream<BIT_ACTV> s_output_golden("s_output_golden");
    maxpool2x2_golden<N_CH, N_IH, N_IW, N_BATCH, BIT_ACTV>(
        s_input_golden, s_output_golden
    );


    // Check results
    auto error_cnt = check_afm<BIT_ACTV, 4>(s_output_golden, s_output, ofm_shape, 1);
    
    if(error_cnt > 0) {
        std::cout << "Test Failed! " << error_cnt << " errors" << std::endl;
        return 1;
    } else {
        std::cout << "Test Passed! " << std::endl;
        return 0;
    }
}
